import cv2
import torch
import pickle
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms


# Pickle functions to save and load dictionaries
def save_obj(obj, name):
    with open(name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


def load_obj(name):
    with open(name + '.pkl', 'rb') as f:
        return pickle.load(f)

# helper function to plot grpahs


def visualize(array):
    "Plot all images in the array of tensors in one row"
    for z in range(0, len(array)):
        temp = array[z]
        if temp.ndim > 3:  # tensor output in the form NCHW
            temp = (torch.argmax(temp, dim=1).squeeze())
        if len(temp.shape) >= 3:
            plt.figure()
            plt.imshow(np.transpose(
                temp.detach().numpy().squeeze(), (1, 2, 0)))
            plt.show()
        else:
            plt.figure()
            plt.imshow(temp.detach().numpy(), cmap='gray')


def create_bounding_box(img, pose, pt_cld_data, intrinsic_matrix,color=(0,0,255)):
    "Create a bounding box around the object"
    # 8 corner points of the ptcld data
    min_x, min_y, min_z = pt_cld_data.min(axis=0)
    max_x, max_y, max_z = pt_cld_data.max(axis=0)
    # print(max_x-min_x,max_y-min_y,max_z-min_z)
    # max_z = 20
    corners_3D = np.array([[max_x, min_y, min_z],
                           [max_x, min_y, max_z],
                           [min_x, min_y, max_z],
                           [min_x, min_y, min_z],
                           [max_x, max_y, min_z],
                           [max_x, max_y, max_z],
                           [min_x, max_y, max_z],
                           [min_x, max_y, min_z]])

    # convert these 8 3D corners to 2D points
    ones = np.ones((corners_3D.shape[0], 1))
    homogenous_coordinate = np.append(corners_3D, ones, axis=1)

    # Perspective Projection to obtain 2D coordinates for masks
    homogenous_2D = intrinsic_matrix @ (pose @ homogenous_coordinate.T)
    coord_2D = homogenous_2D[:2, :] / homogenous_2D[2, :]
    coord_2D = ((np.floor(coord_2D)).T).astype(int)
    # print(coord_2D)
    # Draw lines between these 8 points
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[1]), color, 1)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[3]), color, 1)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[4]), color, 1)
    img = cv2.line(img, tuple(coord_2D[1]), tuple(coord_2D[2]), color, 1)
    img = cv2.line(img, tuple(coord_2D[1]), tuple(coord_2D[5]), color, 1)
    img = cv2.line(img, tuple(coord_2D[2]), tuple(coord_2D[3]), color, 1)
    img = cv2.line(img, tuple(coord_2D[2]), tuple(coord_2D[6]), color, 1)
    img = cv2.line(img, tuple(coord_2D[3]), tuple(coord_2D[7]), color, 1)
    img = cv2.line(img, tuple(coord_2D[4]), tuple(coord_2D[7]), color, 1)
    img = cv2.line(img, tuple(coord_2D[4]), tuple(coord_2D[5]), color, 1)
    img = cv2.line(img, tuple(coord_2D[5]), tuple(coord_2D[6]), color, 1)
    img = cv2.line(img, tuple(coord_2D[6]), tuple(coord_2D[7]), color, 1)

    return img


def set_message(img,position, message):
    "Add a message to the image"
    font = cv2.FONT_HERSHEY_SIMPLEX
    fontScale = 0.8
    color = (145, 24, 241)
    thickness = 1
    img = cv2.putText(img, message, position , font,
                      fontScale, color, thickness, cv2.LINE_AA)
    return img



def create_coor(img, pose, intrinsic_matrix,message=""):
    "Create a bounding box around the object"
    # 8 corner points of the ptcld data

    x_color = (255, 0, 0)
    y_color = (0, 255, 0)
    z_color = (0, 0, 255)
    corners_3D = np.array([[0, 0, 0],
                           [0.05, 0, 0],
                           [0, 0.05, 0],
                           [0, 0, 0.05]  ])

    # convert these 8 3D corners to 2D points
    ones = np.ones((corners_3D.shape[0], 1))
    homogenous_coordinate = np.append(corners_3D, ones, axis=1)

    # Perspective Projection to obtain 2D coordinates for masks
    homogenous_2D = intrinsic_matrix @ (pose @ homogenous_coordinate.T)
    coord_2D = homogenous_2D[:2, :] / homogenous_2D[2, :]
    coord_2D = ((np.floor(coord_2D)).T).astype(int)

    # Draw lines between these 8 points
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[1]), x_color, 1)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[2]), y_color, 1)
    img = cv2.line(img, tuple(coord_2D[0]), tuple(coord_2D[3]), z_color, 1)

    img = set_message(img, tuple(coord_2D[0]), message)
    return img

def ADD_score(pt_cld, true_pose, pred_pose, diameter):
    "Evaluation metric - ADD score"
    pred_pose[0:3, 0:3][np.isnan(pred_pose[0:3, 0:3])] = 1
    pred_pose[:, 3][np.isnan(pred_pose[:, 3])] = 0
    target = pt_cld @ true_pose[0:3, 0:3] + np.array(
        [true_pose[0, 3], true_pose[1, 3], true_pose[2, 3]])
    output = pt_cld @ pred_pose[0:3, 0:3] + np.array(
        [pred_pose[0, 3], pred_pose[1, 3], pred_pose[2, 3]])
    avg_distance = (np.linalg.norm(output - target))/pt_cld.shape[0]
    threshold = diameter * 0.1
    # if avg_distance <= threshold:
    #     return avg_distance
    # else:
    #     return 0
    return avg_distance


if __name__ == "__main__":
    print("This is a helper file. It contains functions to visualize the bounding box and pose on the image.")
    import sys,os
    sys.path.append(os.getcwd())
    from InitPose.models.pose_cnn_fft import PoseCNN 
    from InitPose.lib.dataset.PROPSPoseDataset import PROPSPoseDataset
    from InitPose.lib.options.config import Config
    from torch.utils.data import DataLoader
    import random
    from PIL import Image
    id =0
    
    train_dataset = PROPSPoseDataset("./","train")

    train_dataloader = DataLoader(dataset=train_dataset, batch_size=1,num_workers=4)
    print(f"Dataset sizes: train ({len(train_dataset)})")
    sample_idx = random.randint(0,len(train_dataloader.dataset)-1) 

    data=train_dataloader.dataset[sample_idx]
    img = data['rgb']
    image_data_transposed = img.transpose((1, 2, 0))  # 转置为 (高度, 宽度, 通道数)
    image_rgb =Image.fromarray((image_data_transposed * 255).astype(np.uint8))
    image_rgb.save('output_image.png')
    poses =data['RTs'][id]
    # 将位姿值饶x轴旋转90度
    # rotation = np.array([
    # [1, 0, 0,],
    # [0, 0, -1,],
    # [0, 1, 0,],
    #     ])
    # # print(poses)
    # # # 将位姿矩阵饶x轴旋转90度
    # poses[:3,:3] = rotation @ poses[:3,:3]
    # # poses = np.linalg.inv(poses)
    
    print(poses)
    image_rgb = np.array(image_rgb)

    pt_cld_data= train_dataset.models_pcd[id]
    intrinsic_matrix = train_dataset.cam_intrinsic
    box_imge =create_bounding_box(image_rgb, poses, pt_cld_data, intrinsic_matrix)
    coord_img = create_coor(image_rgb, poses, intrinsic_matrix,message='GT')
    
    
    
    cv2.imwrite('output_image_with_bounding_box.png',box_imge)
    cv2.imwrite('output_image_with_coor.png',coord_img)
    
    
    
    